import csv
import glob
import os

def add_ratio_column():
    """为每个 csv 文件添加一个新列：average_generate_tokens_per_step / total_forward_batch_spec_time_ms"""
    
    print("=== 为每个文件添加比率列 ===\n")
    
    # 获取所有 csv 文件
    csv_files = glob.glob("timings_all_setting_rlhf_*.csv")
    
    for file_path in sorted(csv_files):
        file_name = os.path.basename(file_path)
        print(f"处理文件: {file_name}")
        
        rows = []
        
        # 读取原文件
        with open(file_path, newline='') as f:
            reader = csv.DictReader(f)
            fieldnames = reader.fieldnames + ['time_ratio']
            
            for row in reader:
                if row.get('total_forward_batch_spec_time_ms') and row.get('average_generate_tokens_per_step'):
                    try:
                        time_ms = float(row['total_forward_batch_spec_time_ms'])
                        avg_tokens = float(row['average_generate_tokens_per_step'])
                        # 计算比率：average_generate_tokens_per_step / total_forward_batch_spec_time_ms
                        ratio = avg_tokens / time_ms
                        row['time_ratio'] = ratio
                    except (ValueError, KeyError):
                        row['time_ratio'] = ''
                else:
                    row['time_ratio'] = ''
                
                rows.append(row)
        
        # 写回新文件
        out_path = file_path.replace('.csv', '_with_ratio.csv')
        with open(out_path, 'w', newline='') as f:
            writer = csv.DictWriter(f, fieldnames=fieldnames)
            writer.writeheader()
            writer.writerows(rows)
        
        print(f"  已生成: {out_path}")
    
    print(f"\n总共处理了 {len(csv_files)} 个文件")

if __name__ == "__main__":
    add_ratio_column() 